import sys 
sys.path.append("..") 
from graph import Graph_resnet18
import torch
import json
import time
import os

FRAME_N = 5000
CONFIG_FILE = "config.json"


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
g = Graph_resnet18()

# warm up 
x = torch.randn(1, 3, 224, 224).to(device)
for _ in range(5):
    g(x)

while True:
    print("standing by ...")
    while not os.path.exists(CONFIG_FILE):
        time.sleep(1)
    config = json.load(open("config.json"))
    os.remove(CONFIG_FILE)
    start = config['start']
    end = config['end']
    frame_id = 0

    x = torch.randn(1, 3, 224, 224).to(device)
    if (start <= 1):
        shape = x.shape
    else:
        shape = g(x, start=0, end=start)[0].shape


    x0 = torch.randn(shape).to(device)
    x1 = torch.randn(shape).to(device)
    N = 10

    while True:
        starter, ender = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
        torch.cuda.synchronize()
        print("  block=", config)
        print("{0: >10s} {1: >20s}".format("frame_id", "elapsed_time(ms)"))
        print("-------------------------------------")
        avg = 0
        for i in range(N):
            starter.record()
            g(x0, x1, start, end)
            ender.record()
            torch.cuda.synchronize()
            curr_time = starter.elapsed_time(ender) # 从 starter 到 ender 之间用时,单位为毫秒
            print("{0: >10d} |{1: >15.4f}".format(frame_id, curr_time))
            avg += curr_time
            frame_id += 1

        print("-------------------------------------")
        print("{0: >18s}=  {1: >1.4f}\n\n".format("avg", avg / N))
        if frame_id > FRAME_N:
            break